#!/usr/bin/env python3
"""
视觉识别ROS节点
该节点负责从一个或多个ROS图像话题接收图像，并在收到指令后，
使用火山引擎的Ark多模态大模型（doubao-1.5-vision-pro）进行场景识别。
识别结果（场景描述文本）将以流式方式发布到TTS话题，供语音合成节点使用。
"""
import os
from typing import Any, List, cast
import rospy
from sensor_msgs.msg import Image
from std_msgs.msg import String
from cv_bridge import CvBridge
import cv2
import base64
from volcenginesdkarkruntime import Ark
import threading

class VisualRecognitionNode:
    """
    封装了视觉识别所有功能的ROS节点类。
    - 订阅图像话题和控制指令话题。
    - 使用CV Bridge转换图像格式。
    - 调用多模态大模型API进行分析。
    - 将结果发布到TTS话题。
    """
    def __init__(self):
        """
        初始化视觉识别节点。
        """
        rospy.init_node('visual_recognition', anonymous=True)
        
        # 从ROS参数服务器获取要订阅的图像话题列表
        raw_topics = rospy.get_param('~image_topics', ['visual_info'])
        self.topics: List[str] = [raw_topics] if isinstance(raw_topics, str) else raw_topics
        
        # 从ROS参数服务器获取要使用的视觉大模型名称
        raw_model = rospy.get_param('~model', 'doubao-1.5-vision-pro-250328')
        self.model: str = str(raw_model)
        
        # 初始化CvBridge，用于ROS图像消息和OpenCV图像格式之间的转换
        self.bridge = CvBridge()
        # 初始化缓冲区，用于同步来自不同话题的图像
        self.buffers = {}
        self.expected_count = len(self.topics)
        
        # 初始化TTS文本发布器
        self.pub = rospy.Publisher('tts_text', String, queue_size=10)
        # 初始化火山Ark大模型客户端
        self.client = Ark(api_key='XXXXXXXXXXXXXXXXXXX')
        
        # 订阅指定的图像话题
        for topic in self.topics:
            rospy.Subscriber(topic, Image, self.image_callback)
        
        # 订阅视觉识别指令话题，用于触发识别过程
        rospy.Subscriber('visual_command', String, self.visual_command_callback)
        
        # 状态标志，控制是否处理接收到的图像，初始为False
        self.should_process = False
        
        rospy.loginfo(f"已订阅图像话题: {self.topics}, 使用模型: {self.model}")
        rospy.loginfo("视觉识别节点已启动，等待 'visual_command' 指令...")

    def visual_command_callback(self, msg):
        """
        处理来自'visual_command'话题的指令。
        当收到"true"时，将should_process置为True，以触发一次图像处理。
        """
        if msg.data.lower() == "true":
            rospy.loginfo("收到场景识别指令，准备处理下一帧图像。")
            self.should_process = True
        else:
            rospy.logwarn(f"接收到未知的视觉指令: {msg.data}，已忽略。")

    def image_callback(self, msg: Image):
        """
        图像话题的回调函数。
        如果should_process为False，则忽略图像。
        否则，对图像进行编码，并与其他话题的图像进行同步。
        当收集到所有话题的图像后，启动一个新线程进行处理。
        """
        if not self.should_process:
            return
        
        # 使用时间戳作为图像帧的唯一标识，用于同步
        ts = msg.header.stamp.to_nsec()
        try:
            # 将ROS Image消息转换为OpenCV图像格式
            cv_img = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
            # 将图像编码为JPEG格式，然后进行Base64编码，以符合API要求
            ret, buf = cv2.imencode('.jpg', cv_img)
            if not ret:
                raise ValueError('JPEG编码失败')
            raw_b64 = base64.b64encode(buf.tobytes()).decode('utf-8')
            data_uri = f"data:image/jpeg;base64,{raw_b64}"
        except Exception as e:
            rospy.logerr(f"图像转换或编码失败: {e}")
            return
            
        # 将处理后的图像数据存入对应时间戳的缓冲区
        lst = self.buffers.setdefault(ts, [])
        lst.append(data_uri)
        
        # 如果已集齐所有来源的图像
        if len(lst) >= self.expected_count:
            images = self.buffers.pop(ts)
            # 重置处理状态，确保每次指令只处理一次
            self.should_process = False
            # 在新线程中调用大模型API，避免阻塞回调函数
            threading.Thread(target=self.process_images_stream, args=(images,)).start()

    def process_images_stream(self, images: List[str]):
        """
        调用火山Ark大模型API处理图像，并以流式方式发布结果。
        
        Args:
            images (List[str]): 包含一个或多个Base64编码图像URI的列表。
        """
        # 构建发送给大模型的内容
        content: List[dict] = []
        for img_b64 in images:
            content.append({'type': 'image_url', 'image_url': {'url': img_b64}})
        # 添加文本提示，指导模型如何响应
        content.append({'type': 'text', 'text': '请描述你看到了什么。控制在40字以内'})
        
        rospy.loginfo("图像已集齐，正在调用大模型进行视觉识别...")
        try:
            # 调用大模型的流式API
            resp_stream = cast(Any, self.client.chat.completions.create(
                model=self.model,
                messages=[{'role': 'user', 'content': content}],
                stream=True
            ))

            buffer = ""
            # 处理流式响应的每个数据块
            for chunk in resp_stream:
                chunk_content = chunk.choices[0].delta.content
                if chunk_content:
                    buffer += chunk_content
                    # 当缓冲区中包含句末标点时，认为一句话已完整，进行发布
                    if any(p in buffer for p in ["。", "！", "？", "...", "\n"]):
                        # 找到第一个出现的句末标点并分割
                        punc = next(p for p in ["。", "！", "？", "...", "\n"] if p in buffer)
                        sentence, _, remainder = buffer.partition(punc)
                        sentence_to_publish = (sentence + punc).strip()
                        if sentence_to_publish:
                            self.pub.publish(sentence_to_publish)
                            rospy.loginfo(f"已发布句子: {sentence_to_publish}")
                        buffer = remainder

            # 处理流末尾可能剩余的文本
            final_sentence = buffer.strip()
            if final_sentence:
                self.pub.publish(final_sentence)
                rospy.loginfo(f"已发布最后一个句子: {final_sentence}")
                
        except Exception as e:
            rospy.logerr(f"调用Ark API时发生错误: {e}")
            self.pub.publish("抱歉，我在识别图像时遇到了问题。")

if __name__ == '__main__':
    try:
        node = VisualRecognitionNode()
        rospy.spin()
    except rospy.ROSInterruptException:
        rospy.loginfo("视觉识别节点关闭。")
        pass